#include "opcode.h"
#include "symbol.h"
#include <iostream>
#include <cassert>

using namespace std;

extern SymbolTable symbolTable;
FunctionManager* FunctionManager::instance = 0;
Context* Context::curContext = 0;


Context::Context(FunctionSymbol *sym){
    funcSym = sym;
    nextSlot = 0;
    nextReg = 0;
    dstReg = NODEST;
}

int Context::regToVar(int reg){
    // use the sequential search, it may not harm the performance much
    for(map<int,int>::iterator it = varToReg.begin();it != varToReg.end();it++){
        if(it->second == reg){
            return it->first;
        }
    }

    return -1;
}

OpCodeGen::OpCodeGen(){
    pc = 1;
}

OpCodeGen::~OpCodeGen(){

}

const char* OpCodeGen::getCodeName(int op){
    static const char* codeName[] = {
        "NOP",
        "ADD",
        "SUB",
        "MUL",
        "DIV",
        "MOV",
        "GETLOCAL",
        "SETLOCAL",
        "GETGLOBAL",
        "SETGLOBAL",
        "GOTO",
        "LDCONST",
        "PUSH",
        "POP"
    };
    if(op >= sizeof(codeName)/sizeof(codeName[0])){
        return 0;
    }

    return codeName[op];
}

void OpCodeGen::printInfo(){
    vector<FunctionSymbol*>::iterator it;
    FunctionManager *mngr = FunctionManager::getInstance();
    for(it = mngr->functions.begin();it != mngr->functions.end();it++){
        ops_type &ops = (*it)->ops;
        cout<<((*it)->getName());
        if((*it)->funcName){
            cout<<": LINE "<<(*it)->funcName->defLine;
        }
        cout<<endl;
        for(int i = 1;i <= ops.size();i++){
            Opcode *op = ops[i];
            cout<<"#"<<i<<"  ";
            op->print();
        }
        cout<<endl;
    }
}


void OpCodeGen::visitVariableExpr(VariableExpr& expr){
    Context *ctx = Context::curContext;
    Symbol *sym = symbolTable.getSymbol(expr.name,ctx);


    /// SymbolEnter should enter it
    assert(sym && "symbol hasn't been entered");

    int local = sym->getLocal();
    if(ctx->varToReg.find(local) != ctx->varToReg.end()){
        // if the variable has been allocated, just return it
        result = ctx->varToReg[local];
    }else{
        // otherwise, designate a register for it
        ctx->varToReg[local] = ctx->nextReg;
        result = ctx->nextReg++;
    }

    //assert(resultReg != Symbol::NOLOCAL);
}

void OpCodeGen::visitConstantExpr(ConstantExpr& expr){
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;
    ConstantSymbol  *sym = expr.constSym;

    // load the constant into next available register
    assert((sym->getPoolIndex() != Symbol::NOINDEX) &&
           "the constant must be in constant pool");
    ops[pc++] = new LoadConstOp(ctx->nextReg,sym->getPoolIndex());
    result = ctx->nextReg++;
}

void OpCodeGen::visitBinaryExpr(BinaryExpr& expr){
    int src1,src2,dst;
    Context *ctx = Context::curContext;

    expr.left->accept(this);
    src1 = result;
    expr.right->accept(this);
    src2 = result;

    dst = ctx->nextReg++;

    Opcode* op;
    switch(expr.getOperator()){
    case ADD:
        op = new AddOp(dst,src1,src2);break;
    case SUB:
        op = new SubOp(dst,src1,src2);break;
    case MUL:
        op = new MulOp(dst,src1,src2);break;
    case DIV:
        op = new DivOp(dst,src1,src2);break;
    default:
        assert(0 && "unknown operator");
        break;
    }

    ops_type&   ops = ctx->funcSym->ops;
    ops[pc++] = op;
    result = dst;
}

void OpCodeGen::visitAssignment(AssignmentStatement& stmt){
    int lhs,rhs;
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;

    stmt.lhs->accept(this);
    lhs = result;
    stmt.rhs->accept(this);
    rhs = result;

    ops[pc++] = new MovOp(lhs,rhs);
}

void OpCodeGen::visitParameterList(ParameterListStatement& stmt){
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;
    vector<VariableExpr*>::iterator it;

    // already in a called function, pop all caller-saved parameters to localVar
    for(it = stmt.variables.begin();it != stmt.variables.end();it++){
        // the SymbolEnter treats the parameters as local variables
        // and assigns them the slot index prior to the latter
        // so the handling is somewhat similar to local variables
        VarSymbol *var = (*it)->varSym;
        assert(var && "Parameters in function definition must be variables");

        // allocate a register for the slot, needless to search them in symbol table
        ctx->varToReg[var->getLocal()] = ctx->nextReg;
        ops[pc++] = new PopOp(ctx->nextReg++);
    }
}

void OpCodeGen::visitExprList(ExprListStatement& stmt){
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;
    vector<Expr*>::reverse_iterator it;

    // still within the calling function, push them before CallOp
    for(it = stmt.variables.rbegin();it != stmt.variables.rend();it++){
        // enqueue the variables(expressions) to stack with reverse order
        // and ensure they are all defined
        // TODO: add constant support and replace here by a checker class
        VariableExpr *vexpr = dynamic_cast<VariableExpr*>(*it);
        assert(vexpr && "must be defined and currently only support local variables");
        VarSymbol *sym = vexpr->varSym;
        assert(sym && "VarSymbol must exist");

        ops[pc++] = new PushOp(sym->getLocal());

    }

}

void OpCodeGen::visitFunctionDef(FunctionDefStatement& stmt){
    // record the function entry
    stmt.funcName->funcSym->entryAddr = pc;
    int oldpc = pc;
    Context *oldctx = Context::curContext;

    Context::curContext = stmt.funcName->funcSym->context;
    ops_type&   ops = stmt.funcName->funcSym->ops;
    // new function pc start at 1
    pc = 1;

    if(stmt.paraList){
        stmt.paraList->accept(this);
    }

    // then get the return address and the return slot index
    int retSlot = Context::curContext->nextSlot++;
    ops[pc++] = new PopOp(retSlot);

    // TODO: how to deal with return value, by stack or by register?
    stmt.funcBody->accept(this);
    ops[pc++] = new RetOp(retSlot);


    Context::curContext = oldctx;
    pc = oldpc;

    // record this function
    FunctionManager::getInstance()->functions.push_back(stmt.funcName->funcSym);
}

void OpCodeGen::visitFunctionCall(FunctionCallStatement& stmt){
    int entry = stmt.funcName->funcSym->entryAddr;
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;

    assert(entry && "this function may not be defined before using it");
    stmt.exprList->accept(this);

    // the operand is the address other than register index
    ops[pc++] = new CallOp(entry);
}

void OpCodeGen::visitBlock(Block &blk){
    Scope *oldscp = Scope::curScope;
    Scope::curScope = blk.blockSym->scope;

    vector<Tree*>::iterator it;
    for(it = blk.statements.begin();it != blk.statements.end();it++){
        (*it)->accept(this);
    }

    Scope::curScope = oldscp;
}

void OpCodeGen::visitWhileStatement(WhileStatement &stmt){
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;
    int beginPc,truePc,falsePc,stubPc;
    int condReg;

    beginPc = pc;
    stmt.expr->accept(this);
    stubPc = pc++;
    condReg = result;

    if(stmt.whileBody){
        truePc = pc;
        stmt.whileBody->accept(this);
        ops[pc++] = new GotoOp(beginPc);
    }else{
        truePc = beginPc;
    }
    falsePc = pc;

    ops[stubPc] = new CondGotoOp(condReg,truePc,falsePc);
}

void OpCodeGen::visitIfStatement(IfStatement &stmt){
    Context *ctx = Context::curContext;
    ops_type&   ops = ctx->funcSym->ops;
    int condPc,ifTrueJmpPc;
    int trueLabel,falseLabel,endLabel;
    int condReg;

    stmt.expr->accept(this);
    condReg = result;
    condPc = pc++;
    trueLabel = pc;

    if(stmt.ifBlock){
        stmt.ifBlock->accept(this);
    }
    ifTrueJmpPc = pc++;

    if(stmt.elseBlock){
        falseLabel = pc;
        stmt.elseBlock->accept(this);
        endLabel = pc;
    }else{
        // the else clause is not available
        falseLabel = endLabel = pc;
    }

    ops[condPc] = new CondGotoOp(condReg,trueLabel,falseLabel);
    ops[ifTrueJmpPc] = new GotoOp(endLabel);


}
